#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8

import scipy.misc
import scipy.stats as stats
from math import *


class Binominal():
    '''
        continuity correction 的文章: 什么时候 +0.5, 什么时候-0.5
        http://www.statisticshowto.com/what-is-the-continuity-correction-factor/
    '''
    ThresHold = 10      # 可以用正太分布模拟的阈值

    def __init__(self, n, p, accept_appro=True):
        self.p = p
        self.q = 1 - p
        self.n = n

        self.mean = n * p              # 期望
        self.sd = sqrt(n * p * self.q) # 标准差
        self.cdp = [0] * (n + 1)       # 累积分布概率表: cdp[n] == (小于等于n个发生的概率)
        self.norm_simulate = 0         # 是否使用正态模拟

        if accept_appro:
            # 接受近似值
            if self.n * self.p >= self.ThresHold or self.n * self.q >= self.ThresHold:
                self.norm_simulate = 1      # 是否使用正态模拟
                return

        # 小于阈值/不接受模拟值: 计算累积分布概率
        self._calc_cdp()
        self.norm_simulate = 0

    def _calc_cdp(self):
        ''' 计算二项累积分布概率
            先计算第0项的值, 目的是去掉for中的边界判断.
        '''
        prob_0 = self.q ** self.n
        self.cdp[0] = prob_0

        for idx in range(1, self.n + 1):
            non_cnt = self.n - idx
            prob_i = scipy.misc.comb(self.n, idx) * (self.p ** idx) * (self.q ** non_cnt)
            self.cdp[idx] = prob_i + self.cdp[idx - 1]

        assert self.cdp[self.n] - 1 <= 0.00001      # 累积分布最终值为1

    def z_score(self, n):
        return (n - self.mean) / self.sd

    def below_p(self, cnt):
        ''' 二项分布小于等于cnt的概率: 比如打靶试验击中[0, 20]次的概率 '''
        if self.norm_simulate:
            # 返回正态估计值, 因为是小于等于, 所以continuity correction factor(fudge factor)+0.5
            z_score = self.z_score(cnt + 0.5)
            return stats.norm.cdf(z_score)
        else:
            # 返回准确值
            return self.cdp[cnt]

def test():
    # n 越大, 估计的越好
    n = 1000
    p = 0.5
    less = 500

    b = Binominal(n, p, True)
    print(b.below_p(less))    # 0.588

    b = Binominal(n, p, False)
    print(b.below_p(less))    # 0.588


def main():
    print("start main")
    test()

if __name__ == "__main__":
    main()
